package com.liuqi.openai.data.document;

import com.liuqi.openai.data.document.parser.TextDocumentParser;
import com.liuqi.openai.data.document.source.FileDocumentSource;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Stream;

import static com.liuqi.openai.util.ExceptionUtil.illegalArgument;
import static java.nio.file.Files.isDirectory;
import static java.nio.file.Files.isRegularFile;

/**
 * 文档加载工具类
 *
 * @author liuqi
 * @date 2025/8/7
 **/
public class DocumentLoader {

    /**
     * 文件类型的文档加载器
     */
    public static class FileDocumentLoader {

        public static Document load(Path filePath) {
            return load(filePath, DocumentType.of(filePath.toString()));
        }

        public static Document load(String filePath) {
            return load(Paths.get(filePath));
        }

        public static Document load(String filePath, DocumentType documentType) {
            return load(Paths.get(filePath), documentType);
        }

        public static List<Document> loads(String directoryPath) {
            return loads(Paths.get(directoryPath));
        }

        public static List<Document> loads(Path directoryPath) {
            if (!isDirectory(directoryPath)) {
                throw illegalArgument("%s is not a directory", directoryPath);
            }

            List<Document> documents = new ArrayList<>();

            try (Stream<Path> paths = Files.list(directoryPath)) {
                paths.filter(Files::isRegularFile)
                        .forEach(filePath -> {
                            try {
                                Document document = load(filePath);
                                documents.add(document);
                            } catch (Exception e) {
                                throw illegalArgument(e, "failed to load document from " + filePath);
                            }
                        });
            } catch (IOException e) {
                throw new RuntimeException(e);
            }

            return documents;
        }

        public static Document load(Path filePath, DocumentType documentType) {
            if (!isRegularFile(filePath)) {
                throw illegalArgument("%s is not a file", filePath);
            }
            return DocumentLoader.load(FileDocumentSource.from(filePath), parserFor(documentType));
        }
    }

    /**
     * 加载文档
     *
     * @param source 文档数据源
     * @param parser 文档解析器
     * @return
     */
    public static Document load(DocumentSource source, DocumentParser parser) {
        try (InputStream inputStream = source.inputStream()) {
            Document document = parser.parse(inputStream);
            source.metadata().asMap().forEach((key, value) -> document.metadata().add(key, value));
            return document;
        } catch (Exception e) {
            throw new RuntimeException("Failed to load document", e);
        }
    }

    /**
     * 通过文档类型获取解析器
     *
     * @param type 文档类型
     * @return
     */
    public static DocumentParser parserFor(DocumentType type) {
        switch (type) {
            case TXT:
            case UNKNOWN:
                return new TextDocumentParser(type);
            default:
                throw new RuntimeException(String.format("Cannot find parser for document type '%s'", type));
        }
    }

}
